#include <cuda_runtime.h>
#include <iostream>

// 矩阵乘法的CUDA内核函数
__global__ void matrixMulCUDA(float *C, float *A, float *B, int wA, int wB)
{
    // 每个线程计算一个结果元素
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (row < wA && col < wB) {
        float value = 0;
        for (int k = 0; k < wA; ++k)
        {
            value += A[row * wA + k] * B[k * wB + col];
        }
        C[row * wB + col] = value;
    }
}

int main()
{
    int wA = 3; // 矩阵A的宽度
    int hA = 3; // 矩阵A的高度
    int wB = 3; // 矩阵B的宽度
    int hB = 3; // 矩阵B的高度

    // 主机端内存分配
    float h_A[] = {1, 2, 3, 4, 5, 6, 7, 8, 9};
    float h_B[] = {9, 8, 7, 6, 5, 4, 3, 2, 1};
    float h_C[wA * hB];

    float *d_A, *d_B, *d_C;

    // 设备端内存分配
    cudaMalloc((void **)&d_A, wA * hA * sizeof(float));
    cudaMalloc((void **)&d_B, wB * hB * sizeof(float));
    cudaMalloc((void **)&d_C, wA * hB * sizeof(float));

    // 将数据从主机端复制到设备端
    cudaMemcpy(d_A, h_A, wA * hA * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, wB * hB * sizeof(float), cudaMemcpyHostToDevice);

    // 定义线程块和线程网格的维度
    dim3 threadsPerBlock(16, 16);
    dim3 blocksPerGrid((wB + threadsPerBlock.x - 1) / threadsPerBlock.x,
                       (hA + threadsPerBlock.y - 1) / threadsPerBlock.y);

    // 启动CUDA内核
    matrixMulCUDA<<<blocksPerGrid, threadsPerBlock>>>(d_C, d_A, d_B, wA, wB);

    // 将结果从设备端复制到主机端
    cudaMemcpy(h_C, d_C, wA * hB * sizeof(float), cudaMemcpyDeviceToHost);

    // 输出结果
    for (int i = 0; i < hA; ++i)
    {
        for (int j = 0; j < wB; ++j)
        {
            std::cout << h_C[i * wB + j] << " ";
        }
        std::cout << std::endl;
    }

    // 释放设备端内存
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);

    return 0;
}
